import pymysql
import pandas as pd
from pyecharts import options as opts
from pyecharts.charts import Bar, Line, Pie, Map, Grid, Page, HeatMap
from pyecharts.globals import ThemeType
from pyecharts.commons.utils import JsCode
from datetime import datetime


# 数据库连接配置
def get_db_connection():
    return pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='ukraine_conflict',
        charset='utf8mb4'
    )


# 获取数据
def fetch_data():
    conn = get_db_connection()
    query = "SELECT * FROM russia_ukraine_conflict"
    df = pd.read_sql(query, conn)
    conn.close()
    return df


# 创建时间趋势图
def create_time_trend(df):
    df['event_date'] = pd.to_datetime(df['event_date'])
    daily_events = df.groupby('event_date').size().reset_index(name='count')

    line = (
        Line(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(daily_events['event_date'].dt.strftime('%Y-%m-%d').tolist())
        .add_yaxis("事件数量", daily_events['count'].tolist())
        .set_global_opts(
            title_opts=opts.TitleOpts(title="冲突事件时间趋势"),
            xaxis_opts=opts.AxisOpts(name="日期"),
            yaxis_opts=opts.AxisOpts(name="事件数量"),
            datazoom_opts=[opts.DataZoomOpts()],
        )
    )
    return line


# 创建冲突严重程度和地理分布热力图
def create_conflict_heatmap(df):
    # 处理数据
    df['severity_level'] = pd.qcut(df['fatalities'], q=5, labels=['极低', '低', '中', '高', '极高'])
    severity_counts = df.groupby(['severity_level', 'region']).size().unstack(fill_value=0)

    # 创建热力图
    heatmap = (
        HeatMap(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(severity_counts.columns.tolist())
        .add_yaxis(
            "严重程度分布",
            severity_counts.index.tolist(),
            severity_counts.values.tolist(),
            label_opts=opts.LabelOpts(formatter="{c}")
        )
        .set_global_opts(
            title_opts=opts.TitleOpts(title="冲突严重程度地理分布热力图"),
            xaxis_opts=opts.AxisOpts(
                name="地区",
                axislabel_opts=opts.LabelOpts(rotate=45)
            ),
            yaxis_opts=opts.AxisOpts(name="严重程度"),
            visualmap_opts=opts.VisualMapOpts(
                min_=0,
                max_=severity_counts.values.max(),
                is_piecewise=True,
                pieces=[
                    {"min": 0, "max": 10, "label": "0-10", "color": "#FFE5DB"},
                    {"min": 10, "max": 50, "label": "10-50", "color": "#FFB3A7"},
                    {"min": 50, "max": 100, "label": "50-100", "color": "#FF7F6B"},
                    {"min": 100, "max": 200, "label": "100-200", "color": "#FF4B2F"},
                    {"min": 200, "label": ">200", "color": "#FF0000"}
                ]
            ),
            toolbox_opts=opts.ToolboxOpts(
                feature=opts.ToolBoxFeatureOpts(
                    save_as_image=opts.ToolBoxFeatureSaveAsImageOpts(title="保存为图片"),
                    data_zoom=opts.ToolBoxFeatureDataZoomOpts(title="区域缩放"),
                    restore=opts.ToolBoxFeatureRestoreOpts(title="还原")
                )
            )
        )
    )
    return heatmap


# 创建伤亡人数柱状图
def create_fatalities_bar(df):
    monthly_fatalities = df.groupby(df['event_date'].dt.strftime('%Y-%m'))['fatalities'].sum().reset_index()

    bar = (
        Bar(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(monthly_fatalities['event_date'].tolist())
        .add_yaxis("伤亡人数", monthly_fatalities['fatalities'].tolist())
        .set_global_opts(
            title_opts=opts.TitleOpts(title="月度伤亡人数统计"),
            xaxis_opts=opts.AxisOpts(name="月份"),
            yaxis_opts=opts.AxisOpts(name="伤亡人数"),
            datazoom_opts=[opts.DataZoomOpts()],
        )
    )
    return bar


# 创建地区分布图
def create_region_distribution(df):
    region_counts = df['region'].value_counts()

    bar = (
        Bar(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(region_counts.index.tolist())
        .add_yaxis("事件数量", region_counts.values.tolist())
        .set_global_opts(
            title_opts=opts.TitleOpts(title="地区事件分布"),
            xaxis_opts=opts.AxisOpts(name="地区", axislabel_opts=opts.LabelOpts(rotate=45)),
            yaxis_opts=opts.AxisOpts(name="事件数量"),
        )
    )
    return bar


# 创建主要参与者分析
def create_actor_analysis(df):
    actor1_counts = df['actor1'].value_counts().head(10)

    bar = (
        Bar(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(actor1_counts.index.tolist())
        .add_yaxis("参与事件数", actor1_counts.values.tolist())
        .set_global_opts(
            title_opts=opts.TitleOpts(title="主要参与者分析"),
            xaxis_opts=opts.AxisOpts(name="参与者", axislabel_opts=opts.LabelOpts(rotate=45)),
            yaxis_opts=opts.AxisOpts(name="参与事件数"),
        )
    )
    return bar


# 创建事件类型趋势
def create_event_type_trend(df):
    event_type_monthly = df.groupby([df['event_date'].dt.strftime('%Y-%m'), 'event_type']).size().unstack(fill_value=0)

    line = (
        Line(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(event_type_monthly.index.tolist())
    )

    for event_type in event_type_monthly.columns:
        line.add_yaxis(event_type, event_type_monthly[event_type].tolist())

    line.set_global_opts(
        title_opts=opts.TitleOpts(title="事件类型趋势"),
        xaxis_opts=opts.AxisOpts(name="月份"),
        yaxis_opts=opts.AxisOpts(name="事件数量"),
        datazoom_opts=[opts.DataZoomOpts()],
    )
    return line


# 创建冲突事件类型与伤亡关系图
def create_conflict_analysis(df):
    # 处理数据
    event_fatalities = df.groupby('event_type').agg({
        'fatalities': ['sum', 'count', 'mean']
    }).reset_index()

    event_fatalities.columns = ['event_type', 'total_fatalities', 'event_count', 'avg_fatalities']
    event_fatalities = event_fatalities.sort_values('total_fatalities', ascending=False).head(10)

    # 创建双Y轴图表
    bar = (
        Bar(init_opts=opts.InitOpts(theme=ThemeType.LIGHT))
        .add_xaxis(event_fatalities['event_type'].tolist())
        .add_yaxis(
            "事件数量",
            event_fatalities['event_count'].tolist(),
            yaxis_index=0,
            label_opts=opts.LabelOpts(position="top")
        )
        .extend_axis(
            yaxis=opts.AxisOpts(
                name="平均伤亡人数",
                type_="value",
                position="right"
            )
        )
        .set_global_opts(
            title_opts=opts.TitleOpts(title="主要事件类型与伤亡关系分析"),
            xaxis_opts=opts.AxisOpts(
                name="事件类型",
                axislabel_opts=opts.LabelOpts(rotate=45)
            ),
            yaxis_opts=opts.AxisOpts(
                name="事件数量",
                position="left"
            ),
            datazoom_opts=[
                opts.DataZoomOpts(type_="slider", range_start=0, range_end=100),
                opts.DataZoomOpts(type_="inside", range_start=0, range_end=100)
            ],
            toolbox_opts=opts.ToolboxOpts(
                feature=opts.ToolBoxFeatureOpts(
                    save_as_image=opts.ToolBoxFeatureSaveAsImageOpts(),
                    data_zoom=opts.ToolBoxFeatureDataZoomOpts(),
                    restore=opts.ToolBoxFeatureRestoreOpts()
                )
            )
        )
    )

    # 添加平均伤亡人数折线
    line = (
        Line()
        .add_xaxis(event_fatalities['event_type'].tolist())
        .add_yaxis(
            "平均伤亡人数",
            event_fatalities['avg_fatalities'].round(2).tolist(),
            yaxis_index=1,
            label_opts=opts.LabelOpts(position="top")
        )
    )

    # 组合图表
    bar.overlap(line)
    return bar


def main():
    # 获取数据
    df = fetch_data()

    # 创建页面
    page = Page(layout=Page.SimplePageLayout)
    page.page_title = "乌克兰冲突数据分析大屏"

    # 添加所有图表
    page.add(
        create_time_trend(df),
        create_conflict_analysis(df),  # 新的分析图表
        create_fatalities_bar(df),
        create_region_distribution(df),
        create_actor_analysis(df),
        create_event_type_trend(df)
    )
    # 生成HTML文件
    page.render("ukraine_conflict_dashboard.html")


if __name__ == "__main__":
    main()